# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import getpass
import logging
import subprocess
import sys
import tempfile

from absl import app, flags
import jinja2

from coin.base.datetime_util import iterate_date
from coin2.support.feed_archive.app.symbol_groups_updater import (
    extract_symbol_group_dict_from_toml,
)

FLAGS = flags.FLAGS


class NoJobError(Exception):
  pass


temp_driver_files = []
def generate_script_setting(trading_dates):
  machines = [
      "feed-05.ap-northeast-1.aws",
  ]
  meas = {
      'Futures.Binance.v1',
      'Futures.Okex.v5-swap',
      'Futures.Bybit.v3-linear',
      'Futures.Prex.v1',
  }
  recipes = [
      'bbo_price_change'
  ]
  settings = {}
  for trading_date in trading_dates:
    settings[trading_date] = {}
    for m in machines:
      if (FLAGS.machine is not None) and (m not in FLAGS.machine):
        continue
      settings[trading_date][m] = {}
      for mea in meas:
        if (FLAGS.mea is not None) and (mea not in FLAGS.mea):
          continue
        settings[trading_date][m][mea] = {}
        for recipe in recipes:
          if (FLAGS.recipe is not None) and (recipe not in FLAGS.recipe):
            continue
          settings[trading_date][m][mea][recipe] = {}

          symbol_group_dict = extract_symbol_group_dict_from_toml(mea)
          for group, symbols in symbol_group_dict.items():
            if (FLAGS.group is not None) and (group not in FLAGS.group):
              continue
            filtered = []
            for symbol in symbols:
              if 'PERPETUAL' in symbol:
                filtered.append(symbol)
            settings[trading_date][m][mea][recipe][group] = filtered
  return settings


def generate_slurm_script(
    start_date,
    end_date,
    summary_mailto,
    user=None):
  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  trading_dates = [td for td in iterate_date(start_date.date(), end_date.date())]
  trading_dates = [td.strftime('%Y%m%d') for td in trading_dates]

  settings = generate_script_setting(trading_dates)

  mail_subject = "[feed-dumper] ${SLURM_JOB_ID} summary"

  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)
  template = template_env.get_template(
      'python/coin2/support/feed_dumper/slurm_script.tmpl')
  return template.render(settings=settings,
                         mail_subject=mail_subject,
                         summary_mailto=summary_mailto)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_on_slurm(
    start_date,
    end_date,
    summary_mailto='{user}@prestolabs.io',
    dry_run=False,
    sbatch_args=None,
    verbose=False):
  if summary_mailto is not None:
    if '{user}' in summary_mailto:
      user = 'leon'
      summary_mailto = summary_mailto.format(user=user)
    if verbose:
      sys.stderr.write('Feed dump result mail will be sent to %s.\n' % summary_mailto)

  try:
    script = generate_slurm_script(start_date,
                                   end_date,
                                   summary_mailto,
                                   user=getpass.getuser())
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  tmp_dir = '/remote/iosg/home-8/bot-hft/jenkins/workspace/coin2_feed_dumper/coin/tmp'
  with tempfile.NamedTemporaryFile(mode='w', dir=tmp_dir) as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  for tmp_driver in temp_driver_files:
    tmp_driver.close()
  return job_id


def main(argv):
  try:
    run_on_slurm(start_date=FLAGS.start_date,
                 end_date=FLAGS.end_date,
                 summary_mailto=FLAGS.mailto,
                 dry_run=FLAGS.script_only,
                 sbatch_args=FLAGS.sbatch_args,
                 verbose=True)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  flags.DEFINE_string('start_date', None, 'yyyymmdd')
  flags.DEFINE_string('end_date', None, 'yyyymmdd. Feed at end_date is not checked.')
  flags.DEFINE_string('mailto',
                      '{user}@prestolabs.io',
                      'If specified, it will send a mail after finish.')
  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')
  flags.DEFINE_string(
      'sbatch_args',
      '--wait --hint=nomultithread -n100 --job-name=feeddump --priority=TOP',
      'sbatch arguments')

  flags.DEFINE_list('machine', None, "Separate by comma")

  flags.DEFINE_list('mea', None, "Separate by comma")

  flags.DEFINE_list('recipe', None, "Separate by comma")

  flags.DEFINE_list('group', None, "Separate by comma")

  app.run(main)
